"""
PyTorch Lightning model for segmentation based on camvid_segmentation_multiclass.ipynb
"""

import pytorch_lightning as pl
import segmentation_models_pytorch as smp
import torch
from torch.optim import lr_scheduler
import torch.nn as nn


class SegmentationLightningModel(pl.LightningModule):
    """PyTorch Lightning model for semantic segmentation"""
    
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.save_hyperparameters()
        
        # Build model
        if config.MODEL_NAME == "Segformer":
            self.model = smp.Segformer(
                encoder_name=config.ENCODER_NAME,
                encoder_weights=config.ENCODER_WEIGHTS,
                in_channels=3,
                classes=config.NUM_CLASSES,
                activation=None,
            )
        else:
            model_class = getattr(smp, config.MODEL_NAME)
            self.model = model_class(
                encoder_name=config.ENCODER_NAME,
                encoder_weights=config.ENCODER_WEIGHTS,
                in_channels=3,
                classes=config.NUM_CLASSES,
                activation=None,
            )
        
        # Preprocessing parameters for image normalization
        params = smp.encoders.get_preprocessing_params(config.ENCODER_NAME)
        self.register_buffer("std", torch.tensor(params["std"]).view(1, 3, 1, 1))
        self.register_buffer("mean", torch.tensor(params["mean"]).view(1, 3, 1, 1))
        
        # Loss function for multi-class segmentation
        self.loss_fn = smp.losses.DiceLoss(smp.losses.MULTICLASS_MODE, from_logits=True)
        
        # Step metrics tracking
        self.training_step_outputs = []
        self.validation_step_outputs = []
        self.test_step_outputs = []
    
    def forward(self, image):
        """Forward pass with image normalization"""
        # Normalize image
        image = (image - self.mean) / self.std
        mask = self.model(image)
        return mask
    
    def shared_step(self, batch, stage):
        """Shared step for training, validation, and testing"""
        image, mask = batch
        
        # Ensure that image dimensions are correct
        assert image.ndim == 4  # [batch_size, channels, H, W]
        
        # Ensure the mask is a long (index) tensor
        mask = mask.long()
        
        # Mask shape
        assert mask.ndim == 3  # [batch_size, H, W]
        
        # Predict mask logits
        logits_mask = self.forward(image)
        
        assert logits_mask.shape[1] == self.config.NUM_CLASSES  # [batch_size, number_of_classes, H, W]
        
        # Ensure the logits mask is contiguous
        logits_mask = logits_mask.contiguous()
        
        # Compute loss using multi-class Dice loss
        loss = self.loss_fn(logits_mask, mask)
        
        # Apply softmax to get probabilities for multi-class segmentation
        prob_mask = logits_mask.softmax(dim=1)
        
        # Convert probabilities to predicted class labels
        pred_mask = prob_mask.argmax(dim=1)
        
        # Compute true positives, false positives, false negatives, and true negatives
        tp, fp, fn, tn = smp.metrics.get_stats(
            pred_mask, mask, mode="multiclass", num_classes=self.config.NUM_CLASSES
        )
        
        return {
            "loss": loss,
            "tp": tp,
            "fp": fp,
            "fn": fn,
            "tn": tn,
        }
    
    def shared_epoch_end(self, outputs, stage):
        """Shared epoch end for training, validation, and testing"""
        # Aggregate step metrics
        tp = torch.cat([x["tp"] for x in outputs])
        fp = torch.cat([x["fp"] for x in outputs])
        fn = torch.cat([x["fn"] for x in outputs])
        tn = torch.cat([x["tn"] for x in outputs])
        
        # Per-image IoU and dataset IoU calculations
        per_image_iou = smp.metrics.iou_score(
            tp, fp, fn, tn, reduction="micro-imagewise"
        )
        dataset_iou = smp.metrics.iou_score(tp, fp, fn, tn, reduction="micro")
        
        metrics = {
            f"{stage}_loss": torch.stack([x["loss"] for x in outputs]).mean(),
            f"{stage}_per_image_iou": per_image_iou,
            f"{stage}_dataset_iou": dataset_iou,
        }
        
        self.log_dict(metrics, prog_bar=True)
    
    def training_step(self, batch, batch_idx):
        """Training step"""
        train_loss_info = self.shared_step(batch, "train")
        self.training_step_outputs.append(train_loss_info)
        return train_loss_info
    
    def on_train_epoch_end(self):
        """Training epoch end"""
        self.shared_epoch_end(self.training_step_outputs, "train")
        self.training_step_outputs.clear()
    
    def validation_step(self, batch, batch_idx):
        """Validation step"""
        valid_loss_info = self.shared_step(batch, "valid")
        self.validation_step_outputs.append(valid_loss_info)
        return valid_loss_info
    
    def on_validation_epoch_end(self):
        """Validation epoch end"""
        self.shared_epoch_end(self.validation_step_outputs, "valid")
        self.validation_step_outputs.clear()
    
    def test_step(self, batch, batch_idx):
        """Test step"""
        test_loss_info = self.shared_step(batch, "test")
        self.test_step_outputs.append(test_loss_info)
        return test_loss_info
    
    def on_test_epoch_end(self):
        """Test epoch end"""
        self.shared_epoch_end(self.test_step_outputs, "test")
        self.test_step_outputs.clear()
    
    def configure_optimizers(self):
        """Configure optimizers and schedulers"""
        optimizer = torch.optim.AdamW(
            self.parameters(), 
            lr=self.config.LEARNING_RATE,
            weight_decay=self.config.WEIGHT_DECAY
        )
        
        # Calculate T_max based on dataset size and epochs
        if hasattr(self.config, 'TRAIN_LOADER_LENGTH'):
            T_max = self.config.NUM_EPOCHS * self.config.TRAIN_LOADER_LENGTH
        else:
            T_max = self.config.NUM_EPOCHS * 100  # Default estimate
        
        scheduler = lr_scheduler.CosineAnnealingLR(
            optimizer, 
            T_max=T_max, 
            eta_min=1e-6
        )
        
        return {
            "optimizer": optimizer,
            "lr_scheduler": {
                "scheduler": scheduler,
                "interval": "step",
                "frequency": 1,
            },
        }